package com.agent.platform.mq;

import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.json.JSONUtil;
import com.agent.platform.config.RabbitMQConfig;
import com.agent.platform.entity.AgentInfo;
import com.agent.platform.entity.TaskInfo;
import com.agent.platform.mapper.AgentInfoMapper;
import com.agent.platform.mapper.TaskInfoMapper;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.rabbitmq.client.Channel;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.rabbit.annotation.RabbitListener;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.HashMap;
import java.util.Map;

/**
 * 任务消息消费者
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class TaskConsumer {
    
    private final TaskInfoMapper taskInfoMapper;
    private final AgentInfoMapper agentInfoMapper;
    private final TaskProducer taskProducer;
    
    private static final int MAX_RETRY_COUNT = 3;
    
    /**
     * 消费任务消息
     */
    @RabbitListener(queues = RabbitMQConfig.TASK_QUEUE)
    public void consumeTask(TaskMessage message, Message mqMessage, Channel channel) {
        log.info("收到任务消息: taskId={}, retryCount={}", 
                message.getTaskId(), message.getRetryCount());
        
        try {
            // 查询任务
            TaskInfo task = getTaskById(message.getTaskId());
            if (task == null) {
                log.warn("任务不存在: taskId={}", message.getTaskId());
                channel.basicAck(mqMessage.getMessageProperties().getDeliveryTag(), false);
                return;
            }
            
            // 查找可用智能体
            AgentInfo agent = findAvailableAgent(message.getTaskType());
            if (agent == null) {
                handleNoAgentAvailable(message, mqMessage, channel);
                return;
            }
            
            // 分配任务
            assignTaskToAgent(task, agent);
            
            // 通知智能体
            notifyAgent(agent, task);
            
            // 确认消息
            channel.basicAck(mqMessage.getMessageProperties().getDeliveryTag(), false);
            log.info("任务分配成功: taskId={}, agentId={}", task.getTaskId(), agent.getAgentId());
            
        } catch (Exception e) {
            log.error("处理任务消息失败: taskId={}", message.getTaskId(), e);
            handleError(message, mqMessage, channel);
        }
    }
    
    /**
     * 处理死信队列消息
     */
    @RabbitListener(queues = RabbitMQConfig.TASK_DLX_QUEUE)
    public void consumeDlxTask(TaskMessage message, Message mqMessage, Channel channel) {
        log.warn("任务进入死信队列: taskId={}, retryCount={}", 
                message.getTaskId(), message.getRetryCount());
        
        try {
            // 更新任务状态为失败
            TaskInfo task = getTaskById(message.getTaskId());
            if (task != null && "PENDING".equals(task.getStatus())) {
                task.setStatus("FAILED");
                task.setErrorMessage("任务分配失败，无可用智能体");
                task.setEndTime(LocalDateTime.now());
                taskInfoMapper.updateById(task);
            }
            
            channel.basicAck(mqMessage.getMessageProperties().getDeliveryTag(), false);
        } catch (Exception e) {
            log.error("处理死信消息失败: taskId={}", message.getTaskId(), e);
            try {
                channel.basicNack(mqMessage.getMessageProperties().getDeliveryTag(), false, false);
            } catch (Exception ex) {
                log.error("拒绝消息失败", ex);
            }
        }
    }
    
    private TaskInfo getTaskById(String taskId) {
        LambdaQueryWrapper<TaskInfo> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(TaskInfo::getTaskId, taskId);
        return taskInfoMapper.selectOne(wrapper);
    }
    
    private AgentInfo findAvailableAgent(String taskType) {
        LambdaQueryWrapper<AgentInfo> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(AgentInfo::getStatus, "ONLINE")
               .eq(AgentInfo::getAgentType, taskType)
               .orderByAsc(AgentInfo::getHeartbeatTime)
               .last("LIMIT 1");
        return agentInfoMapper.selectOne(wrapper);
    }
    
    private void assignTaskToAgent(TaskInfo task, AgentInfo agent) {
        task.setAssignedAgentId(agent.getAgentId());
        task.setStatus("RUNNING");
        task.setStartTime(LocalDateTime.now());
        taskInfoMapper.updateById(task);
    }
    
    private void notifyAgent(AgentInfo agent, TaskInfo task) {
        try {
            String url = agent.getEndpoint() + "/task";
            Map<String, Object> data = new HashMap<>();
            data.put("taskId", task.getTaskId());
            data.put("taskType", task.getTaskType());
            data.put("inputData", task.getInputData());
            
            HttpResponse response = HttpRequest.post(url)
                    .body(JSONUtil.toJsonStr(data))
                    .contentType("application/json")
                    .timeout(5000)
                    .execute();
            
            if (!response.isOk()) {
                log.warn("通知智能体失败: agentId={}, status={}", 
                        agent.getAgentId(), response.getStatus());
            }
        } catch (Exception e) {
            log.error("通知智能体异常: agentId={}", agent.getAgentId(), e);
        }
    }
    
    private void handleNoAgentAvailable(TaskMessage message, Message mqMessage, Channel channel) {
        try {
            if (message.getRetryCount() < MAX_RETRY_COUNT) {
                // 重试
                message.setRetryCount(message.getRetryCount() + 1);
                channel.basicNack(mqMessage.getMessageProperties().getDeliveryTag(), false, false);
                
                // 延迟重新发送
                Thread.sleep(5000);
                taskProducer.sendTask(message);
                log.info("任务重新入队: taskId={}, retryCount={}", 
                        message.getTaskId(), message.getRetryCount());
            } else {
                // 超过重试次数，拒绝消息进入死信队列
                channel.basicNack(mqMessage.getMessageProperties().getDeliveryTag(), false, false);
                log.warn("任务超过最大重试次数: taskId={}", message.getTaskId());
            }
        } catch (Exception e) {
            log.error("处理无可用智能体失败", e);
        }
    }
    
    private void handleError(TaskMessage message, Message mqMessage, Channel channel) {
        try {
            channel.basicNack(mqMessage.getMessageProperties().getDeliveryTag(), false, true);
        } catch (Exception e) {
            log.error("拒绝消息失败", e);
        }
    }
}
